package com.shanhy.demo.controller;

import io.github.lnyocly.ai4j.listener.SseListener;
import io.github.lnyocly.ai4j.platform.deepseek.chat.DeepSeekChatService;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletion;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletionResponse;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatMessage;
import io.github.lnyocly.ai4j.service.IChatService;
import io.github.lnyocly.ai4j.service.PlatformType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.util.concurrent.Executors;

@Slf4j
@RequestMapping("/ai")
@RestController
public class DeepSeekController {

    @Autowired
    private DeepSeekChatService deepSeekChatService;

    @RequestMapping("/chat")
    public String chat(String question) throws Exception {
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model("deepseek-chat")
                .message(ChatMessage.withUser(question))
                .build();
        ChatCompletionResponse chatResponse = deepSeekChatService.chatCompletion(chatCompletion);
        String content = chatResponse.getChoices().get(0).getMessage().getContent();
        return content;
    }

    @RequestMapping("/chatStream")
    public void chatStream(String question, HttpServletResponse response) throws Exception {
        response.setCharacterEncoding("UTF-8");
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model("deepseek-chat")
                .message(ChatMessage.withUser(question))
                .build();
        final PrintWriter writer = response.getWriter();
        SseListener sseListener = new SseListener() {
            @Override
            protected void send() {
                log.info(this.getCurrStr());
                writer.write(this.getCurrStr());
                writer.flush();
            }
        };
        deepSeekChatService.chatCompletionStream(chatCompletion, sseListener);
        writer.close();
    }

    @GetMapping("/chatStream2")
    public ResponseBodyEmitter getChatMessageStream(@RequestParam String question) {
        final ResponseBodyEmitter emitter = new ResponseBodyEmitter();

        // 创建请求参数
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model("deepseek-chat")
                .message(ChatMessage.withUser(question))
                .build();

        Executors.newSingleThreadExecutor().submit(() -> {
            try {
                SseListener sseListener = new SseListener() {
                    @Override
                    protected void send() {
                        try {
                            emitter.send(this.getCurrStr());
                            log.info(this.getCurrStr());
                        } catch (Exception e) {
                            emitter.completeWithError(e);
                        }
                    }
                };

                // 发送流式数据
                deepSeekChatService.chatCompletionStream(chatCompletion, sseListener);

                // 完成后关闭连接
                emitter.complete();
            } catch (Exception e) {
                log.error("execute getChatMessageStream error", e);
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }
}
